/*
 * ========================================
 * 虚拟内存管理系统 (Virtual Memory Management)
 * ========================================
 * 
 * 这个文件实现了RISC-V Sv39页表管理系统，包括：
 * - 页表创建和销毁
 * - 页表项设置和查询
 * - 地址转换
 * - 虚拟内存映射
 * 
 * RISC-V Sv39页表结构：
 * - 3级页表，每级9位索引
 * - 页大小：4KB
 * - 虚拟地址：39位
 * - 物理地址：56位
 */

#include "../include/types.h"
#include "../include/memlayout.h"
#include "../include/riscv.h"
#include "../include/printf.h"
#include "../include/memory.h"

// ========================================
// 全局变量定义
// ========================================

// 内核页表
// 这是整个系统的根页表，所有地址空间都基于它
static pagetable_t kernel_pagetable = 0;

// 页表统计信息
static uint64 page_table_count = 0;      // 页表数量
static uint64 page_table_entry_count = 0; // 页表项数量

// ========================================
// 内部函数声明
// ========================================

// 创建页表
static pagetable_t pagetable_create(void);

// 销毁页表（暂时未使用）
// static void pagetable_destroy(pagetable_t pagetable);

// 查找页表项（内部实现）
static pte_t* walk_internal(pagetable_t pagetable, uint64 va, int alloc);

// 映射页面（内部实现）
static int mappages_internal(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm);

// 取消映射（内部实现）
static void unmappages_internal(pagetable_t pagetable, uint64 va, uint64 size);

// 检查地址是否有效
static int is_valid_virtual_address(uint64 va);
static int is_valid_physical_address(uint64 pa);

// 内存设置函数
static void* vm_memset(void *dst, int c, uint64 n);

// ========================================
// 内部函数实现
// ========================================

/*
 * 创建新的页表
 * @return 页表指针，失败返回0
 */
static pagetable_t pagetable_create(void)
{
    // 分配一个页面作为页表
    pagetable_t pagetable = (pagetable_t)kalloc();
    if (pagetable == 0) {
        printf("错误：无法分配页表页面\n");
        return 0;
    }
    
    // 清零页表
    vm_memset(pagetable, 0, PGSIZE);
    page_table_count++;
    
    return pagetable;
}

/*
 * 销毁页表（暂时未使用）
 * @param pagetable 页表指针
 */
/*
static void pagetable_destroy(pagetable_t pagetable)
{
    if (pagetable == 0) {
        return;
    }
    
    // 递归释放所有页表页面
    // 这里简化处理，只释放根页表
    kfree(pagetable);
    page_table_count--;
}
*/

/*
 * 查找页表项
 * @param pagetable 页表指针
 * @param va 虚拟地址
 * @param alloc 如果页表不存在是否分配
 * @return 页表项指针
 */
static pte_t* walk_internal(pagetable_t pagetable, uint64 va, int alloc)
{
    if (pagetable == 0) {
        return 0;
    }
    
    // 检查虚拟地址是否有效
    if (!is_valid_virtual_address(va)) {
        return 0;
    }
    
    // 提取各级页号
    uint64 vpn2 = VPN2(va);
    uint64 vpn1 = VPN1(va);
    uint64 vpn0 = VPN0(va);
    
    // 第2级页表
    pte_t *pte2 = &pagetable[vpn2];
    if (!PTE_VALID(*pte2)) {
        if (!alloc) {
            return 0;
        }
        
        // 分配第2级页表
        uint64 pa = (uint64)kalloc();
        if (pa == 0) {
            return 0;
        }
        
        *pte2 = PA2PTE(pa) | PTE_V;
        page_table_entry_count++;
    }
    
    // 第1级页表
    pagetable_t pagetable1 = (pagetable_t)PTE2PA(*pte2);
    pte_t *pte1 = &pagetable1[vpn1];
    if (!PTE_VALID(*pte1)) {
        if (!alloc) {
            return 0;
        }
        
        // 分配第1级页表
        uint64 pa = (uint64)kalloc();
        if (pa == 0) {
            return 0;
        }
        
        *pte1 = PA2PTE(pa) | PTE_V;
        page_table_entry_count++;
    }
    
    // 第0级页表
    pagetable_t pagetable0 = (pagetable_t)PTE2PA(*pte1);
    pte_t *pte0 = &pagetable0[vpn0];
    
    return pte0;
}

/*
 * 映射页面
 * @param pagetable 页表指针
 * @param va 虚拟地址
 * @param size 大小
 * @param pa 物理地址
 * @param perm 权限
 * @return 成功返回0，失败返回-1
 */
static int mappages_internal(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm)
{
    if (pagetable == 0) {
        return -1;
    }
    
    // 检查地址对齐
    if ((va & PGMASK) != 0 || (pa & PGMASK) != 0) {
        printf("错误：地址未页对齐 va=0x%p pa=0x%p\n", (void*)va, (void*)pa);
        return -1;
    }
    
    // 检查物理地址是否有效
    if (!is_valid_physical_address(pa)) {
        printf("错误：无效的物理地址 0x%p\n", (void*)pa);
        return -1;
    }
    
    // 逐页映射
    for (uint64 a = va; a < va + size; a += PGSIZE, pa += PGSIZE) {
        pte_t *pte = walk_internal(pagetable, a, 1);
        if (pte == 0) {
            printf("错误：无法创建页表项 va=0x%p\n", (void*)a);
            return -1;
        }
        
        if (PTE_VALID(*pte)) {
            printf("错误：页面已映射 va=0x%p\n", (void*)a);
            return -1;
        }
        
        *pte = PA2PTE(pa) | perm | PTE_V;
        page_table_entry_count++;
    }
    
    return 0;
}

/*
 * 取消映射
 * @param pagetable 页表指针
 * @param va 虚拟地址
 * @param size 大小
 */
static void unmappages_internal(pagetable_t pagetable, uint64 va, uint64 size)
{
    if (pagetable == 0) {
        return;
    }
    
    // 逐页取消映射
    for (uint64 a = va; a < va + size; a += PGSIZE) {
        pte_t *pte = walk_internal(pagetable, a, 0);
        if (pte != 0 && PTE_VALID(*pte)) {
            *pte = 0;
            page_table_entry_count--;
        }
    }
}

/*
 * 检查虚拟地址是否有效
 * @param va 虚拟地址
 * @return 1表示有效，0表示无效
 */
static int is_valid_virtual_address(uint64 va)
{
    // 检查地址是否在有效范围内
    if (va >= MAXVA) {
        return 0;
    }
    
    return 1;
}

/*
 * 检查物理地址是否有效
 * @param pa 物理地址
 * @return 1表示有效，0表示无效
 */
static int is_valid_physical_address(uint64 pa)
{
    // 检查地址是否在有效范围内
    // 允许内核内存范围 (0x80000000 - 0x88000000)
    // 也允许设备内存范围 (0x10000000 - 0x10001000)
    if ((pa >= KERNBASE && pa < PHYSTOP) || 
        (pa >= UART0 && pa < UART0 + PGSIZE)) {
        // 检查地址是否页对齐
        if ((pa & PGMASK) != 0) {
            return 0;
        }
        return 1;
    }
    
    return 0;
}

// ========================================
// 公共接口函数
// ========================================

/*
 * 初始化虚拟内存系统
 */
void vm_init(void)
{
    printf("初始化虚拟内存系统...\n");
    
    // 创建内核页表
    kernel_pagetable = pagetable_create();
    if (kernel_pagetable == 0) {
        printf("错误：无法创建内核页表\n");
        return;
    }
    
    // 映射内核代码和数据段
    // 这里简化处理，直接映射整个内核空间
    uint64 kernel_size = 0x100000;  // 1MB内核空间
    int result = mappages_internal(kernel_pagetable, KERNBASE, kernel_size, KERNBASE, 
                         PTE_R | PTE_W | PTE_X);
    if (result != 0) {
        printf("错误：无法映射内核空间\n");
        return;
    }
    
    // 映射UART设备
    result = mappages_internal(kernel_pagetable, UART0, PGSIZE, UART0, 
                     PTE_R | PTE_W);
    if (result != 0) {
        printf("错误：无法映射UART设备\n");
        return;
    }
    
    printf("虚拟内存系统初始化完成！\n");
    printf("  页表数量：%d\n", (int)page_table_count);
    printf("  页表项数量：%d\n", (int)page_table_entry_count);
    printf("\n");
}

/*
 * 启用虚拟内存
 */
void vm_enable(void)
{
    if (kernel_pagetable == 0) {
        printf("错误：内核页表未初始化\n");
        return;
    }
    
    // 设置SATP寄存器
    uint64 satp = SATP_MODE_SV39 | (((uint64)kernel_pagetable) >> PGSHIFT);
    
    // 写入SATP寄存器
    asm volatile("csrw satp, %0" : : "r"(satp));
    
    // 刷新TLB
    sfence_vma();
    
    printf("虚拟内存已启用！\n");
    printf("  SATP = 0x%p\n", (void*)satp);
    printf("  页表根地址 = 0x%p\n", (void*)kernel_pagetable);
    printf("\n");
}

/*
 * 获取内核页表
 * @return 内核页表指针
 */
pagetable_t get_kernel_pagetable(void)
{
    return kernel_pagetable;
}

/*
 * 打印页表统计信息
 */
void vm_print_stats(void)
{
    printf("=== 虚拟内存统计 ===\n");
    printf("页表数量：%d\n", (int)page_table_count);
    printf("页表项数量：%d\n", (int)page_table_entry_count);
    printf("==================\n\n");
}

/*
 * 测试页表功能
 */
void vm_test(void)
{
    printf("=== 页表功能测试 ===\n");
    
    // 测试地址转换
    uint64 test_va = 0x80000000;
    pte_t *pte = walk_internal(kernel_pagetable, test_va, 0);
    if (pte != 0 && PTE_VALID(*pte)) {
        uint64 pa = PTE2PA(*pte);
        printf("地址转换测试：\n");
        printf("  虚拟地址：0x%p\n", (void*)test_va);
        printf("  物理地址：0x%p\n", (void*)pa);
        printf("  页表项：0x%p\n", (void*)*pte);
        printf("  读权限：%s\n", PTE_READABLE(*pte) ? "是" : "否");
        printf("  写权限：%s\n", PTE_WRITABLE(*pte) ? "是" : "否");
        printf("  执行权限：%s\n", PTE_EXECUTABLE(*pte) ? "是" : "否");
    } else {
        printf("地址转换测试失败：无法找到页表项\n");
    }
    
    printf("==================\n\n");
}

/*
 * 内存设置函数
 * @param dst 目标地址
 * @param c 要设置的字节值
 * @param n 字节数
 * @return 目标地址
 */
static void* vm_memset(void *dst, int c, uint64 n)
{
    char *d = (char*)dst;
    
    for (uint64 i = 0; i < n; i++) {
        d[i] = c;
    }
    
    return dst;
}

// ========================================
// 公共接口函数实现
// ========================================

/*
 * 查找页表项（公共接口）
 */
pte_t* walk(pagetable_t pagetable, uint64 va, int alloc)
{
    return walk_internal(pagetable, va, alloc);
}

/*
 * 映射页面（公共接口）
 */
int mappages(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm)
{
    return mappages_internal(pagetable, va, size, pa, perm);
}

/*
 * 取消映射（公共接口）
 */
void unmappages(pagetable_t pagetable, uint64 va, uint64 size)
{
    unmappages_internal(pagetable, va, size);
}
